import os
import requests
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent
from dotenv import load_dotenv

# 加载环境变量（存储 API 密钥）
load_dotenv()

class MCPWebSearchTool:
    def __init__(self):
        self.tavily_api_key = os.getenv("TAVILY_API_KEY")
        self.base_url = "https://api.tavily.com/search"

    def search(self, query: str, max_results: int = 5) -> str:
        """调用 Tavily API 执行网络搜索"""
        if not self.tavily_api_key:
            return "错误：未配置 Tavily API 密钥"

        params = {
            "api_key": self.tavily_api_key,
            "query": query,
            "max_results": max_results,
            "include_answer": True
        }

        try:
            response = requests.get(self.base_url, params=params)
            response.raise_for_status()
            data = response.json()

            # 格式化搜索结果
            if "results" in data:
                output = []
                for result in data["results"]:
                    title = result.get("title", "无标题")
                    url = result.get("url", "无链接")
                    content = result.get("content", "无内容")
                    output.append(f"## {title}\n链接：{url}\n摘要：{content}\n")
                return "\n".join(output)
            else:
                return "未找到相关结果"
        except Exception as e:
            return f"搜索失败：{str(e)}"

class MCPWebSearchTool_Mock:
    """网络搜索工具 - MCP协议实现"""
    
    def search(self, query: str, max_results: int = 3) -> dict:
        """执行网络搜索（简化版）"""
        # 这里可以使用SerpAPI或其他搜索API
        # 简化实现返回模拟数据
        return {
            'query': query,
            'results': [
                f'搜索结果1: 关于"{query}"的信息...',
                f'搜索结果2: 更多关于"{query}"的内容...'
            ],
            'success': True
        }
